import broadscope_bailian
from broadscope_bailian import ChatQaMessage
from config import AI_CONFIG

# 文档https://help.aliyun.com/document_detail/2587502.html?spm=a2c4g.2587501.0.0.5a493d1fKejiRz


class BailianBot:

    def __init__(self):
        config = AI_CONFIG["bailian"]
        self.access_key_id = config["access_key_id"]
        self.access_key_secret = config["access_key_secret"]
        self.agent_key = config["agent_key"]
        self.app_id = config["app_id"]

    def generate(self, prompt, session_id=None, chat_history=None):
        client = broadscope_bailian.AccessTokenClient(access_key_id=self.access_key_id,
                                                      access_key_secret=self.access_key_secret)
        token = client.get_token()
        broadscope_bailian.api_key = token

        if isinstance(chat_history, list):
            new_history = [ChatQaMessage(m) for m in chat_history]
            response = broadscope_bailian.Completions().call(app_id=self.app_id, prompt=prompt, history=new_history)
        elif session_id is not None:
            response = broadscope_bailian.Completions().call(app_id=self.app_id, prompt=prompt, session_id=session_id)
        else:
            response = broadscope_bailian.Completions().call(app_id=self.app_id, prompt=prompt)

        return response

